/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include "ccl_buffer_manager.h"
#include "log.h"
#include "config_log.h"
#include "externalinput_pub.h"
#include "workflow_pub.h"
#include "adapter_rts_common.h"

namespace hccl {
CCLBufferManager::CCLBufferManager()
    :inCCLbuffer_(DeviceMem()), outCCLbuffer_(DeviceMem()), winExpBuffer_(DeviceMem()),
    inCCLbufferSize_(0), outCCLbufferSize_(0), winExpBufferSize_(0),
    inAlltoAllvParaBuffer_(DeviceMem()), outAlltoAllvParaBuffer_(DeviceMem())
{
}

CCLBufferManager::~CCLBufferManager()
{
    ReleaseCommCCLbuffer();
    ReleaseAlltoAllvParaBuffer();
    ReleaseCommAIVbuffer();
}

HcclResult CCLBufferManager::CreateCCLbuffer(u64 size, DeviceMem &buffer)
{
    CHK_PRT_RET(!size, HCCL_INFO("[CCLBufferManager][CreateCCLbuffer]buffer size is zero. not need to malloc memory"),
        HCCL_SUCCESS);

    CHK_PRT_RET((size > ULONG_MAX),
        HCCL_ERROR("[CCLBufferManager][CreateCCLbuffer]buffer size is greater than %llu", ULONG_MAX), HCCL_E_PARA);

    CHK_RET(DeviceMem::alloc(buffer, size));
    HCCL_INFO("[CreateCCLbuffer] buffer ptr[%p], size[%llu]", buffer.ptr(), buffer.size());
    CHK_PRT_RET(size && !buffer, HCCL_ERROR("[CCLBufferManager][CreateCCLbuffer]Create ccl buffer size[%llu] fail,"
        "please check environmental variable HCCL_BUFFSIZE.", size), HCCL_E_PTR);
    HCCL_RUN_INFO("[HCCL_TRACE][CreateCCLbuffer]Create ccl buffer success. buffer ptr[%p], size[%llu]",
        buffer.ptr(), buffer.size());
    return HCCL_SUCCESS;
}

HcclResult CCLBufferManager::CreateCommCCLbuffer()
{
    if (inCCLbufferSize_ == 0) {
        inCCLbufferSize_ = GetExternalInputCCLBuffSize();
    }
    if (outCCLbufferSize_ == 0) {
        outCCLbufferSize_ = GetExternalInputCCLBuffSize();
    }
    if (winExpBufferSize_ == 0) {
        winExpBufferSize_ = EXP_BUFFER_SIZE;
    }
 
    if (cclBuffer_.ptr() == nullptr) {
        u64 totalSize = inCCLbufferSize_ + outCCLbufferSize_ + winExpBufferSize_;
        CHK_RET(CreateCCLbuffer(totalSize, cclBuffer_));
        CHK_RET(hrtMemSet(cclBuffer_.ptr(), totalSize, totalSize));
    }
 
    if (inCCLbuffer_.ptr() == nullptr) {
        inCCLbuffer_ = DeviceMem::create(cclBuffer_.ptr(), inCCLbufferSize_);
    }
 
    if (outCCLbuffer_.ptr() == nullptr) {
        outCCLbuffer_ = DeviceMem::create(static_cast<u8 *>(cclBuffer_.ptr()) + inCCLbufferSize_, outCCLbufferSize_);
    }
 
    if (winExpBuffer_.ptr() == nullptr) {
        winExpBuffer_ = DeviceMem::create(static_cast<u8 *>(cclBuffer_.ptr()) + inCCLbufferSize_ + outCCLbufferSize_, 
            winExpBufferSize_);
    }
 
    return HCCL_SUCCESS;
}

HcclResult CCLBufferManager::CleanCCLbuffer()
{
    if (inCCLbuffer_.ptr() != nullptr) {
        CHK_RET(hrtMemSet(inCCLbuffer_.ptr(), inCCLbuffer_.size(), inCCLbuffer_.size()));
        HCCL_INFO("[CleanCCLbuffer] clean input buffer, ptr[%p], size[%llu]", inCCLbuffer_.ptr(), inCCLbuffer_.size());
    }

    if (outCCLbuffer_.ptr() != nullptr) {
        CHK_RET(hrtMemSet(outCCLbuffer_.ptr(), outCCLbuffer_.size(), outCCLbuffer_.size()));
        HCCL_INFO("[CleanCCLbuffer] clean output buffer, ptr[%p], size[%llu]",
            outCCLbuffer_.ptr(), outCCLbuffer_.size());
    }
    return HCCL_SUCCESS;
}

HcclResult CCLBufferManager::CleanAIVbuffer(void *bufferPtr)
{
    constexpr u32 MEM_SIZE_1M = 1024 * 1024;
    bool isAivOpsExc = UNLIKELY(GetDebugConfig() & HCCL_AIV_OPS_EXC);
    s64 moreMemory = isAivOpsExc ? MEM_SIZE_1M : 0;
    // 将aiv的bufferPtr空间置于0
    if (bufferPtr != nullptr) {
        CHK_RET(hrtMemSet(bufferPtr, AIV_FLAG_SIZE + moreMemory, AIV_FLAG_SIZE + moreMemory));
        HCCL_INFO("[CleanAIVbuffer] clean aiv buffer, ptr[%p], size[%llu]", bufferPtr, AIV_FLAG_SIZE + moreMemory);
    }
    return HCCL_SUCCESS;
}

HcclResult CCLBufferManager::CreateCommAIVbuffer(bool useOpbaseFlag)
{
    bool isAivOpsExc = UNLIKELY(GetDebugConfig() & HCCL_AIV_OPS_EXC);
    constexpr u32 MEM_SIZE_1M = 1024 * 1024;
    size_t offset = 2 * MEM_SIZE_1M - sizeof(int32_t);
    if (useOpbaseFlag) {
        if (inAivOpbaseBuffer_.ptr() == nullptr) {
            CHK_RET(CreateCCLbuffer(AIV_DATA_SIZE + (isAivOpsExc ? 1 * MEM_SIZE_1M : 0), inAivOpbaseBuffer_));
            CHK_RET(CleanAIVbuffer(static_cast<u8 *>(inAivOpbaseBuffer_.ptr()) + (AIV_DATA_SIZE - AIV_FLAG_SIZE)));
        }
        if (outAivOpbaseBuffer_.ptr() == nullptr) {
            CHK_RET(CreateCCLbuffer(AIV_FLAG_SIZE + (isAivOpsExc ? 1 * MEM_SIZE_1M : 0), outAivOpbaseBuffer_));
            CHK_RET(CleanAIVbuffer(outAivOpbaseBuffer_.ptr()));
            // AivOutBuffer的第2m中最后int32位置存放环境变量
            int32_t *envVarAddr = reinterpret_cast<int32_t *>(reinterpret_cast<uintptr_t>(outAivOpbaseBuffer_.ptr()) + offset);
            HCCL_INFO("[CreateCommAIVbuffer] outAivOpbaseBuffer addr is [%p]", outAivOpbaseBuffer_.ptr());
            int envAivOps[1] = {isAivOpsExc ? 1 : 0};
            CHK_RET(hrtMemSyncCopy(
                envVarAddr, sizeof(int32_t),
                envAivOps, sizeof(int32_t),
                HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_HOST_TO_DEVICE));
            HCCL_RUN_INFO("[HCCL_TRACE][CreateCommAIVbuffer] OpbaseMode");
        }
    } else {
        if (inAivOffloadbuffer_.ptr() == nullptr) {
            CHK_RET(CreateCCLbuffer(AIV_DATA_SIZE + (isAivOpsExc ? 1 * MEM_SIZE_1M : 0), inAivOffloadbuffer_));
            CHK_RET(CleanAIVbuffer(static_cast<u8 *>(inAivOffloadbuffer_.ptr()) + (AIV_DATA_SIZE - AIV_FLAG_SIZE)));
        }
        if (outAivOffloadbuffer_.ptr() == nullptr) {
            CHK_RET(CreateCCLbuffer(AIV_FLAG_SIZE + (isAivOpsExc ? 1 * MEM_SIZE_1M : 0), outAivOffloadbuffer_));
            CHK_RET(CleanAIVbuffer(outAivOffloadbuffer_.ptr()));
            int32_t *envVarAddr = reinterpret_cast<int32_t *>(reinterpret_cast<uintptr_t>(outAivOffloadbuffer_.ptr()) + offset);
            HCCL_INFO("[CreateCommAIVbuffer] outAivOpbaseBuffer addr is [%p]", outAivOpbaseBuffer_.ptr());
            int envAivOps[1] = {isAivOpsExc ? 1 : 0};
            CHK_RET(hrtMemSyncCopy(
                envVarAddr, sizeof(int32_t),
                envAivOps, sizeof(int32_t),
                HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_HOST_TO_DEVICE));
            HCCL_RUN_INFO("[HCCL_TRACE][CreateCommAIVbuffer] OffloadMode");
        }
    }
    return HCCL_SUCCESS;
}

HcclResult CCLBufferManager::CreateCommInfoAIVbuffer()
{
    if (aivCommInfoBuffer_.ptr() == nullptr) {
        CHK_RET(CreateCCLbuffer(AIV_COMM_INFO_SIZE, aivCommInfoBuffer_));
    }
    return HCCL_SUCCESS;
}

HcclResult CCLBufferManager::ReleaseCommCCLbuffer()
{
    if ((cclBuffer_.ptr() == nullptr) && (inCCLbuffer_.ptr() == nullptr) && (outCCLbuffer_.ptr() == nullptr) &&
        (winExpBuffer_.ptr() == nullptr)) {
        HCCL_RUN_INFO("[HCCL_TRACE][ReleaseCCLbuffer]CCLBuffer is null, no need to release.");
        return HCCL_SUCCESS;
    }
 
    if (cclBuffer_.ptr() != nullptr){
        HCCL_RUN_INFO("[HCCL_TRACE][ReleaseCCLbuffer]Release cclBuffer. buffer ptr[%p], size[%llu]",
            cclBuffer_.ptr(), cclBuffer_.size());
        cclBuffer_.free();
    }
 
    if (inCCLbuffer_.ptr() != nullptr){
        HCCL_RUN_INFO("[HCCL_TRACE][ReleaseCCLbuffer]Release incclBuffer. buffer ptr[%p], size[%llu]",
        inCCLbuffer_.ptr(), inCCLbuffer_.size());
        inCCLbuffer_.free();
    }
 
    if (outCCLbuffer_.ptr() != nullptr ){
        HCCL_RUN_INFO("[HCCL_TRACE][ReleaseCCLbuffer]Release outcclBuffer. buffer ptr[%p], size[%llu]",
        outCCLbuffer_.ptr(), outCCLbuffer_.size());
        outCCLbuffer_.free();
    }

    if (winExpBuffer_.ptr() != nullptr ){
        HCCL_RUN_INFO("[HCCL_TRACE][ReleaseCCLbuffer]Release expcclBuffer. buffer ptr[%p], size[%llu]",
        winExpBuffer_.ptr(), winExpBuffer_.size());
        winExpBuffer_.free();
    }
 
    if ((cclBuffer_.ptr() == nullptr) && (inCCLbuffer_.ptr() == nullptr) && (outCCLbuffer_.ptr() == nullptr) &&
        (winExpBuffer_.ptr() == nullptr)) {
        HCCL_RUN_INFO("[HCCL_TRACE][ReleaseCCLbuffer]Release CCLbuffer success.");
    }
    return HCCL_SUCCESS;
}

HcclResult CCLBufferManager::ReleaseCommAIVbuffer()
{
    HCCL_RUN_INFO("[HCCL_TRACE][ReleaseAIVbuffer]Release inAivOpbaseBuffer. buffer ptr[%p], size[%llu]",
        inAivOpbaseBuffer_.ptr(), inAivOpbaseBuffer_.size());
    inAivOpbaseBuffer_.free();
    HCCL_RUN_INFO("[HCCL_TRACE][ReleaseAIVbuffer]Release outAivOpbaseBuffer. buffer ptr[%p], size[%llu]",
        outAivOpbaseBuffer_.ptr(), outAivOpbaseBuffer_.size());
    outAivOpbaseBuffer_.free();
    HCCL_RUN_INFO("[HCCL_TRACE][ReleaseAIVbuffer]Release inAivOffloadbuffer. buffer ptr[%p], size[%llu]",
        inAivOffloadbuffer_.ptr(), inAivOffloadbuffer_.size());
    inAivOffloadbuffer_.free();
    HCCL_RUN_INFO("[HCCL_TRACE][ReleaseAIVbuffer]Release outAivOffloadbuffer. buffer ptr[%p], size[%llu]",
        outAivOffloadbuffer_.ptr(), outAivOffloadbuffer_.size());
    outAivOffloadbuffer_.free();
    HCCL_RUN_INFO("[HCCL_TRACE][ReleaseAIVbuffer]Release aivCommInfoBuffer. buffer ptr[%p], size[%llu]",
        aivCommInfoBuffer_.ptr(), aivCommInfoBuffer_.size());
    aivCommInfoBuffer_.free();
    if (inAivOpbaseBuffer_.ptr() == nullptr && outAivOpbaseBuffer_.ptr() == nullptr &&
        inAivOffloadbuffer_.ptr() == nullptr && outAivOffloadbuffer_.ptr() == nullptr &&
        aivCommInfoBuffer_.ptr() == nullptr) {
        HCCL_RUN_INFO("[HCCL_TRACE][ReleaseAIVbuffer]Release AIV buffer success.");
    }
    return HCCL_SUCCESS;
}

HcclResult CCLBufferManager::ClearCommAIVbuffer()
{
    if (inAivOpbaseBuffer_.ptr() != nullptr) {
        CHK_RET(CleanAIVbuffer(static_cast<u8 *>(inAivOpbaseBuffer_.ptr()) + (AIV_DATA_SIZE - AIV_FLAG_SIZE)));
    }
    if (outAivOpbaseBuffer_.ptr() != nullptr) {
        CHK_RET(CleanAIVbuffer(outAivOpbaseBuffer_.ptr()));
    }
    if (inAivOffloadbuffer_.ptr() != nullptr) {
        CHK_RET(CleanAIVbuffer(static_cast<u8 *>(inAivOffloadbuffer_.ptr()) + (AIV_DATA_SIZE - AIV_FLAG_SIZE)));
    }
    if (outAivOffloadbuffer_.ptr() != nullptr) {
        CHK_RET(CleanAIVbuffer(outAivOffloadbuffer_.ptr()));
    }
    return HCCL_SUCCESS;
}

DeviceMem& CCLBufferManager::GetInAivOpbaseBuffer()
{
    return inAivOpbaseBuffer_;
}

DeviceMem& CCLBufferManager::GetOutAivOpbaseBuffer()
{
    return outAivOpbaseBuffer_;
}

DeviceMem& CCLBufferManager::GetInAivOffloadbuffer()
{
    return inAivOffloadbuffer_;
}

DeviceMem& CCLBufferManager::GetOutAivOffloadbuffer()
{
    return outAivOffloadbuffer_;
}

DeviceMem& CCLBufferManager::GetCommCCLBuffer()
{
    return cclBuffer_;
}

HcclResult CCLBufferManager::InitCCLbuffer(u64 inCCLbufferSize, u64 outCCLbufferSize)
{
    inCCLbufferSize_ = inCCLbufferSize;
    outCCLbufferSize_ = outCCLbufferSize;
    return HCCL_SUCCESS;
}

void* CCLBufferManager::GetCCLbufferAddr(const DeviceMem &buffer)
{
    if (!buffer) {
        return nullptr;
    } else {
        return static_cast<void *>(reinterpret_cast<u8 *>(buffer.ptr()));
    }
}

DeviceMem& CCLBufferManager::GetInCCLbuffer()
{
    return inCCLbuffer_;
}

DeviceMem& CCLBufferManager::GetCommExpBuffer()
{
    return winExpBuffer_;
}

DeviceMem& CCLBufferManager::GetAivCommInfoBuffer()
{
    return aivCommInfoBuffer_;
}

HcclResult CCLBufferManager::GetInCCLbuffer(void* &buffer, u64 &size)
{
    buffer = GetCCLbufferAddr(inCCLbuffer_);
    size = inCCLbufferSize_;
    return HCCL_SUCCESS;
}

u64 CCLBufferManager::GetInCCLbufferSize()
{
    return inCCLbufferSize_;
}

DeviceMem& CCLBufferManager::GetOutCCLbuffer()
{
    return outCCLbuffer_;
}

HcclResult CCLBufferManager::GetOutCCLbuffer(void* &buffer, u64 &size)
{
    buffer = GetCCLbufferAddr(outCCLbuffer_);
    size = outCCLbufferSize_;
    return HCCL_SUCCESS;
}

u64 CCLBufferManager::GetOutCCLbufferSize()
{
    return outCCLbufferSize_;
}

u64 CCLBufferManager::GetExpBufferSize()
{
    return winExpBufferSize_;
}

DeviceMem CCLBufferManager::GetCommRegMem(const DeviceMem &mem, MemAttr memAttr, bool aivMode)
{
    u64 commMemSize = 0;
    if ((GetWorkflowMode() == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) && (!aivMode)) {
        // 单算子模式时，仅在第一次集合通信时创建子通信域，注册通信内存。需要将整个CCLbuffer注册进通信域。
        if (memAttr == MemAttr::IN_CCL_BUFFER) {
            commMemSize = inCCLbufferSize_;
        } else if (memAttr == MemAttr::OUT_CCL_BUFFER) {
            commMemSize = outCCLbufferSize_;
        }
    } else {
        commMemSize = mem.size();
    }
    DeviceMem commMem = DeviceMem::create(mem.ptr(), commMemSize);
    return commMem;
}

HcclResult CCLBufferManager::InitAlltoAllvParaBuffer(u64 inBufferSize, u64 outBufferSize)
{
    CHK_RET(CreateCCLbuffer(inBufferSize, inAlltoAllvParaBuffer_));
    CHK_RET(CreateCCLbuffer(outBufferSize, outAlltoAllvParaBuffer_));
    return HCCL_SUCCESS;
}

DeviceMem& CCLBufferManager::GetInAlltoAllvParaBuffer()
{
    return inAlltoAllvParaBuffer_;
}

DeviceMem& CCLBufferManager::GetOutAlltoAllvParaBuffer()
{
    return outAlltoAllvParaBuffer_;
}

void CCLBufferManager::ReleaseAlltoAllvParaBuffer()
{
    inAlltoAllvParaBuffer_.free();
    outAlltoAllvParaBuffer_.free();
}

HcclResult CCLBufferManager::GetIndependentOpCCLbuffer(void* &buffer, uint64_t &size)
{
    if (inCCLbufferSize_ == 0) {
        HCCL_ERROR("[CCLBufferManager][GetIndependentOpCCLbuffer] inCCLbufferSize is 0");
        return HCCL_E_INTERNAL;
    }
    buffer = GetCCLbufferAddr(inCCLbuffer_);
    if (inCCLbuffer_.ptr() == nullptr) {
        CHK_RET(CreateCCLbuffer(inCCLbufferSize_, inCCLbuffer_));
        buffer = GetCCLbufferAddr(inCCLbuffer_);
    }
    size = inCCLbufferSize_;
    return HCCL_SUCCESS;
}
} // namespace hccl